from collections import defaultdict, deque
from typing import Callable, Deque, List, Optional

import h5py
import numpy as np
import pandas as pd

from coin.support.data_replay.reader.interface import (
    RawData,
    IRawDataReader)


class IntervalFeedHdf5Reader(IRawDataReader):
  def __init__(self,
               start_timestamp: int,
               end_timestamp: int,
               symbols: Optional[List[str]],
               columns: Optional[List[str]],
               files: List[str]):
    self._record_cache = self._load_data(
        start_timestamp=start_timestamp,
        end_timestamp=end_timestamp,
        symbols=symbols,
        columns=columns,
        files=files
    )

  def _load_data(self, *, start_timestamp, end_timestamp, symbols, columns, files
  ) -> Deque[RawData]:
    data = defaultdict(dict)
    for file in files:
      symbol_indexes = []
      filtered_symbols = []
      with h5py.File(file, 'r') as f:
        timestamp_universe = f['timestamp']
        symbol_universe = f['universe'][:].astype(str)
        symbols = symbols if symbols is not None else symbol_universe
        for symbol in symbols:
          if symbol in symbol_universe:
            symbol_indexes.append(np.where(symbol_universe == symbol)[0][0])
            filtered_symbols.append(symbol)
        column_universe = f.keys()
        columns = columns if columns is not None else f.keys()
        filtered_columns = [column for column in columns if column in column_universe]

        for column in filtered_columns:
          dataset = f[column]
          assert len(symbol_indexes) == len(filtered_symbols), (
              len(symbol_indexes), len(filtered_symbols))
          for timestamp_idx, value in enumerate(dataset[:]):
            timestamp = timestamp_universe[timestamp_idx]
            if timestamp < start_timestamp:
              continue
            if timestamp > end_timestamp:
              break
            for symbol_idx, filtered_symbol in zip(symbol_indexes, filtered_symbols):
              ts_entry = data[timestamp]
              symbol = symbol_universe[symbol_idx]
              if symbol not in ts_entry:
                ts_entry[filtered_symbol] = {}
              symbol_entry = ts_entry[filtered_symbol]
              symbol_entry.update({column: value[symbol_idx]})
    processed = []
    for timestamp, elem in data.items():
      timestamp = int(timestamp)
      df = pd.DataFrame(elem)
      df['timestamp'] = timestamp
      processed.append(RawData(timestamp=timestamp, value=df))
    return deque(sorted(processed, key=lambda x: x.timestamp))

  def peek(self) -> Optional[RawData]:
    value = self._record_cache[0] if len(self._record_cache) > 0 else None
    return value

  def read(self) -> Optional[RawData]:
    value = self._record_cache.popleft() if len(self._record_cache) > 0 else None
    return value


if __name__ == '__main__':
  import datetime
  from coin.base.datetime_util import to_datetime, to_timestamp_int
  files = ['/remote/iosg/data-2/buckets/feed.derived.interval_h5/coin/main/PT1M/Futures.Binance.v1/20230723/Futures.Binance.v1--ohlc.h5']
  start_timestamp = to_timestamp_int(datetime.datetime(2023, 7, 23))
  end_timestamp = to_timestamp_int(datetime.datetime(2023, 7, 24))
  symbols = ['STMX-USDT.PERPETUAL',]
  columns = ['OPEN_MID', 'CLOSE_MID']
  reader = IntervalFeedHdf5Reader(
      start_timestamp, end_timestamp, symbols, columns, files)
  prev_ts = 0
  while True:
    record = reader.read()
    if record is None:
      break
    curr_ts = record.timestamp
    assert curr_ts >= prev_ts, (prev_ts, curr_ts)
    prev_ts = curr_ts
    record.value['time'] = record.value['timestamp'].apply(lambda x: to_datetime(int(x)))
    print(record.value)
